import torch
import torchvision.transforms as T
# from transformers import AutoImageProcessor, AutoModel
from PIL import Image
import torch.nn as nn

def calculate_similarity(img1_path, img2_path, img_h=336, img_w=224):
    device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
    # print(f"device is {device}")
    transform = T.Compose([
        T.Resize((img_h, img_w)),        # 调整尺寸到目标大小
        T.ToTensor(),                    # 转换为张量
        T.Normalize(                     # 使用DinoV2的标准归一化参数
            mean=(0.485, 0.456, 0.406), 
            std=(0.229, 0.224, 0.225)
        ),
    ])
    # print(f"transform is {transform}")
    model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
    model.eval()  # 设置为评估模式
    # print(f"model is {model}")
    model.to(device)  # 发送到 GPU 上
    image1 = Image.open(img1_path)
    with torch.no_grad():
        inputs1 = transform(image1).unsqueeze(0).to(device)
        outputs1 = model(inputs1)
        image_features1 = outputs1
        # image_features1 = outputs1.last_hidden_state
        # image_features1 = image_features1.mean(dim=1)
        print(f"outputs1 is {image_features1.shape}")
    image2 = Image.open(img2_path)
    with torch.no_grad():
        inputs2 = transform(image2).unsqueeze(0).to(device)
        outputs2 = model(inputs2)
        image_features2 = outputs2
        # image_features2 = outputs2.last_hidden_state
        # image_features2 = image_features2.mean(dim=1)

    cos = nn.CosineSimilarity(dim=0)
    sim = cos(image_features1[0],image_features2[0]).item()
    sim = (sim+1)/2
    print('Similarity:', sim)



if __name__ == '__main__':
    # print("PyTorch 版本:", torch.__version__)
    # print("CUDA 可用:", torch.cuda.is_available())
    # print("CUDA 版本:", torch.version.cuda)  # 关键：查看 PyTorch 编译时使用的 CUDA 版本
    # print("cuDNN 版本:", torch.backends.cudnn.version())
    # print("GPU 设备数:", torch.cuda.device_count())
    # if torch.cuda.is_available():
    #     print("当前 GPU:", torch.cuda.get_device_name(0))
    calculate_similarity("./test1.jpg","./test2.jpg")
    calculate_similarity("./test2.jpg","./test3.jpg")
